import numpy as np
from tqdm import tqdm

from yan_new import CvFo
import torch
import pandas as pd


def eval_egg():
    voc = pd.read_pickle("voc_data.pandas_pickle")
    net = CvFo(len(voc), 32, 8, 1, "egg")
    net.load_egg()
    net.eval()
    loss_func = torch.nn.CrossEntropyLoss()
    data_set = pd.read_pickle("train_data.pandas_pickle")[:1000]
    batch_size = 1
    loss_list = []
    for i in tqdm(range(0, len(data_set), batch_size)):
        j = i + batch_size
        one_data = data_set[i:j]
        two_data = torch.Tensor(one_data).int()
        out = net(two_data[:, :-1])
        loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]).long())
        loss = loss.item()
        loss_list.append(loss)
    sort_loss = np.argsort(loss_list)
    sort_loss_sub = []
    for i in range(0, len(sort_loss), len(sort_loss) // 5):
        j = i + len(sort_loss) // 5
        one = sort_loss[i:j]
        sort_loss_sub.append(one.tolist())
    if  len(sort_loss_sub)!=5:
        sort_loss_sub[-2]+=sort_loss_sub[-1]
        sort_loss_sub=sort_loss_sub[:-1]
    return [[data_set[j]  for j in i]for i in sort_loss_sub]



def eval_dif():
    voc = pd.read_pickle("voc_data.pandas_pickle")
    net = CvFo(len(voc), 32, 8, 3, "differentiation")

    loss_func = torch.nn.CrossEntropyLoss()
    data_set = pd.read_pickle("train_data.pandas_pickle")
    batch_size = 12
    bar = tqdm(len(100 * list(range(0, len(data_set), batch_size))))

    for i in range(0, len(data_set), batch_size):
        j = i + batch_size
        one_data = data_set[i:j]
        two_data = torch.Tensor(one_data).int()
        out = net(two_data[:, :-1])
        loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]).long())
        bar.set_description("loss___{:.5f}".format(loss.item()))
        bar.update(1)


if __name__ == '__main__':
    eval_egg()

